from typing import Dict, List

class CitationFormatter:
    """引用格式化工具"""
    
    def __init__(self):
        self.venue_map = {
            'iclr': 'ICLR',
            'nips': 'NeurIPS',
            'neurips': 'NeurIPS',
            'icml': 'ICML',
            'aaai': 'AAAI',
            'cvpr': 'CVPR',
            'acl': 'ACL',
            'emnlp': 'EMNLP',
            'eccv': 'ECCV',
            'ijcai': 'IJCAI',
            'kdd': 'KDD',
            'journal': 'Journal'
        }

    def format_inline_citation(self, citation: Dict) -> str:
        """格式化行内引用为<sup>number</sup>，多个引用以空格分隔"""
        if isinstance(citation, list):
            return ' '.join([f"<sup>{cite['number']}</sup>" for cite in citation])
        return f"<sup>{citation['number']}</sup>"

    def format_apa_reference(self, citation: Dict) -> str:
        """Format citation as [number] Title. Venue, Year, chunk_id"""
        title = self._clean_title(citation['title'])
        venue = self._format_venue(citation['venue'])
        number = citation.get('number', '')
        chunk_info = citation.get('chunk_id', '').split('_')[-1] if citation.get('chunk_id') else ''
        
        ref = f"[{number}] "
        ref += f"{title}. "
        ref += f"{venue}, {citation['year']}"
        if chunk_info:
            ref += f", chunk {chunk_info}"
            
        return ref
    def _clean_title(self, title: str) -> str:
        """清理标题文本"""
        return title.strip().rstrip('.')
    
    def _format_venue(self, venue: str) -> str:
        """格式化会议名称"""
        venue = venue.lower()
        return self.venue_map.get(venue, venue)
